{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:02:08.728321Z",
     "start_time": "2024-03-05T12:02:08.712213Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "data = pd.read_csv(\"E:/baselines/ZZ500.csv\")\n",
    "code = \"zz500\"\n",
    "# data = data.drop(columns=['Adj Close'])\n",
    "# data = data.drop(columns=['prev_close', 'adjfactor'])\n",
    "# data.rename(columns={'Date':'date', 'Open':'open', 'High':'high', 'Low':'low', 'Close':'close', 'Volume':'volume', 'company':'tic'}, inplace=True)\n",
    "data.rename(columns={'dt':'date', 'kdcode':'tic'}, inplace=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:12:33.793063Z",
     "start_time": "2024-03-05T12:12:33.668708Z"
    }
   },
   "id": "f051da00c6e53883",
   "execution_count": 61
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "              tic       date  close   open   high    low     turnover  \\\n0       000009.SZ 2018-01-02   7.32   7.27   7.33   7.23  140497530.0   \n1       000009.SZ 2018-01-03   7.42   7.32   7.44   7.30  170636551.0   \n2       000009.SZ 2018-01-04   7.64   7.52   7.70   7.51  358089946.0   \n3       000009.SZ 2018-01-05   7.55   7.65   7.65   7.51  196259451.0   \n4       000009.SZ 2018-01-08   7.52   7.55   7.63   7.46  178511079.0   \n...           ...        ...    ...    ...    ...    ...          ...   \n169007  603883.SH 2023-12-25  29.37  29.02  29.58  28.80  101007153.0   \n169008  603883.SH 2023-12-26  29.05  29.26  29.49  28.82   66158144.0   \n169009  603883.SH 2023-12-27  30.08  28.94  30.40  28.80  206474442.0   \n169010  603883.SH 2023-12-28  28.98  30.15  30.20  28.44  264520627.0   \n169011  603883.SH 2023-12-29  29.89  28.97  30.00  28.80  204020452.0   \n\n            volume  \n0       19269993.0  \n1       23092097.0  \n2       46992707.0  \n3       25990264.0  \n4       23703431.0  \n...            ...  \n169007   3457490.0  \n169008   2272587.0  \n169009   6939076.0  \n169010   9053474.0  \n169011   6878338.0  \n\n[169012 rows x 8 columns]",
      "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>tic</th>\n      <th>date</th>\n      <th>close</th>\n      <th>open</th>\n      <th>high</th>\n      <th>low</th>\n      <th>turnover</th>\n      <th>volume</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>000009.SZ</td>\n      <td>2018-01-02</td>\n      <td>7.32</td>\n      <td>7.27</td>\n      <td>7.33</td>\n      <td>7.23</td>\n      <td>140497530.0</td>\n      <td>19269993.0</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>000009.SZ</td>\n      <td>2018-01-03</td>\n      <td>7.42</td>\n      <td>7.32</td>\n      <td>7.44</td>\n      <td>7.30</td>\n      <td>170636551.0</td>\n      <td>23092097.0</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>000009.SZ</td>\n      <td>2018-01-04</td>\n      <td>7.64</td>\n      <td>7.52</td>\n      <td>7.70</td>\n      <td>7.51</td>\n      <td>358089946.0</td>\n      <td>46992707.0</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>000009.SZ</td>\n      <td>2018-01-05</td>\n      <td>7.55</td>\n      <td>7.65</td>\n      <td>7.65</td>\n      <td>7.51</td>\n      <td>196259451.0</td>\n      <td>25990264.0</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>000009.SZ</td>\n      <td>2018-01-08</td>\n      <td>7.52</td>\n      <td>7.55</td>\n      <td>7.63</td>\n      <td>7.46</td>\n      <td>178511079.0</td>\n      <td>23703431.0</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>169007</th>\n      <td>603883.SH</td>\n      <td>2023-12-25</td>\n      <td>29.37</td>\n      <td>29.02</td>\n      <td>29.58</td>\n      <td>28.80</td>\n      <td>101007153.0</td>\n      <td>3457490.0</td>\n    </tr>\n    <tr>\n      <th>169008</th>\n      <td>603883.SH</td>\n      <td>2023-12-26</td>\n      <td>29.05</td>\n      <td>29.26</td>\n      <td>29.49</td>\n      <td>28.82</td>\n      <td>66158144.0</td>\n      <td>2272587.0</td>\n    </tr>\n    <tr>\n      <th>169009</th>\n      <td>603883.SH</td>\n      <td>2023-12-27</td>\n      <td>30.08</td>\n      <td>28.94</td>\n      <td>30.40</td>\n      <td>28.80</td>\n      <td>206474442.0</td>\n      <td>6939076.0</td>\n    </tr>\n    <tr>\n      <th>169010</th>\n      <td>603883.SH</td>\n      <td>2023-12-28</td>\n      <td>28.98</td>\n      <td>30.15</td>\n      <td>30.20</td>\n      <td>28.44</td>\n      <td>264520627.0</td>\n      <td>9053474.0</td>\n    </tr>\n    <tr>\n      <th>169011</th>\n      <td>603883.SH</td>\n      <td>2023-12-29</td>\n      <td>29.89</td>\n      <td>28.97</td>\n      <td>30.00</td>\n      <td>28.80</td>\n      <td>204020452.0</td>\n      <td>6878338.0</td>\n    </tr>\n  </tbody>\n</table>\n<p>169012 rows × 8 columns</p>\n</div>"
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data['date'] = pd.to_datetime(data['date'])\n",
    "grouped_df = data.groupby('tic')\n",
    "l = data['date'].nunique()\n",
    "def filter_groups(group):\n",
    "    return len(group) == l\n",
    "df = grouped_df.filter(filter_groups)\n",
    "df"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:12:34.716076Z",
     "start_time": "2024-03-05T12:12:34.611944Z"
    }
   },
   "id": "f65f7717cdb26322",
   "execution_count": 62
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "116"
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stock_num = df['tic'].nunique()\n",
    "stock_num"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:12:38.793162Z",
     "start_time": "2024-03-05T12:12:38.770158Z"
    }
   },
   "id": "502cbfc82fc2ed6f",
   "execution_count": 63
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "\"\\ntic_sw_mapping = df[['tic', 'sw_kdcode_1']]\\nunique_mapping = tic_sw_mapping.groupby('tic').first().to_dict()['sw_kdcode_1']\\nunique_mapping\\n\""
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stock_unique = df['tic'].unique()\n",
    "'''\n",
    "tic_sw_mapping = df[['tic', 'sw_kdcode_1']]\n",
    "unique_mapping = tic_sw_mapping.groupby('tic').first().to_dict()['sw_kdcode_1']\n",
    "unique_mapping\n",
    "'''"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:12:48.055910Z",
     "start_time": "2024-03-05T12:12:48.042149Z"
    }
   },
   "id": "838112ce4febc0e4",
   "execution_count": 64
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "1457"
     },
     "execution_count": 65,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df['date'].nunique()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:12:51.195024Z",
     "start_time": "2024-03-05T12:12:51.175229Z"
    }
   },
   "id": "a90115201da06910",
   "execution_count": 65
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "1215"
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "'''\n",
    "test_idx的下标\n",
    "'''\n",
    "p = df['date'].unique()\n",
    "np.where(p == \"2023-01-03\")[0][0]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:12:51.723199Z",
     "start_time": "2024-03-05T12:12:51.702025Z"
    }
   },
   "id": "b95e131a9b2f35a1",
   "execution_count": 66
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 116/116 [00:01<00:00, 94.84it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "scaler = MinMaxScaler()\n",
    "# cols_to_normalize = ['close', 'open', 'high', 'low', 'volume']\n",
    "cols_to_normalize = ['close', 'open', 'high', 'low', 'volume', 'turnover']\n",
    "\n",
    "stocks = []\n",
    "\n",
    "for stock in tqdm(df['tic'].unique()):\n",
    "    tmp = df[df['tic'] == stock].reset_index(drop=True)\n",
    "    tmp[cols_to_normalize] = scaler.fit_transform(tmp[cols_to_normalize])\n",
    "    tmp = tmp.drop(columns=['date', 'tic'])\n",
    "    # tmp = tmp.drop(columns=['date', 'tic', 'sw_kdcode_1'])\n",
    "    stocks.append(tmp.values)\n",
    "\n",
    "stocks = np.array(stocks)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:13:12.780764Z",
     "start_time": "2024-03-05T12:13:11.528341Z"
    }
   },
   "id": "785fa9bcef85598d",
   "execution_count": 67
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "df['ret_1'] = df['close'] / df['close'].shift(1)\n",
    "df['ret_1'] = df['ret_1'].fillna(1)\n",
    "df.loc[df['date'] == '2018-01-02', 'ret_1'] = 1\n",
    "df['ret_1'] = df['ret_1'] - 1"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:13:13.768802Z",
     "start_time": "2024-03-05T12:13:13.750298Z"
    }
   },
   "id": "4f86738e67e88a41",
   "execution_count": 68
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 116/116 [00:00<00:00, 121.84it/s]\n"
     ]
    }
   ],
   "source": [
    "ret = []\n",
    "\n",
    "for stock in tqdm(df['tic'].unique()):\n",
    "    tmp = df[df['tic'] == stock].reset_index(drop=True)\n",
    "    ret.append(tmp['ret_1'].values)\n",
    "\n",
    "ret = np.array(ret)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:13:16.571813Z",
     "start_time": "2024-03-05T12:13:15.605634Z"
    }
   },
   "id": "a244e1a5d22cb39c",
   "execution_count": 69
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "(116, 1457, 6)"
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stocks.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:13:17.152225Z",
     "start_time": "2024-03-05T12:13:17.141613Z"
    }
   },
   "id": "a984a990463d800b",
   "execution_count": 70
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "(116, 1457)"
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ret.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:13:18.727143Z",
     "start_time": "2024-03-05T12:13:18.719759Z"
    }
   },
   "id": "d7cd07c1558c208f",
   "execution_count": 71
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "array([[0.14572666, 0.14574257, 0.14129159, 0.15429282, 0.0210878 ,\n        0.0374774 ],\n       [0.14966522, 0.14772277, 0.14559687, 0.15719618, 0.02642707,\n        0.04722418],\n       [0.15833005, 0.15564356, 0.15577299, 0.16590626, 0.05963536,\n        0.10817332],\n       ...,\n       [0.31547853, 0.31683168, 0.31076321, 0.33181253, 0.01440102,\n        0.01091871],\n       [0.32099252, 0.31762376, 0.31741683, 0.33471589, 0.02350612,\n        0.02174417],\n       [0.31981095, 0.32435644, 0.31819961, 0.3405226 , 0.02339496,\n        0.02166309]])"
     },
     "execution_count": 73,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stocks[0]"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:13:46.194430Z",
     "start_time": "2024-03-05T12:13:46.190432Z"
    }
   },
   "id": "5f8cf8d4590ddb70",
   "execution_count": 73
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "np.save(f'{code}/stocks_data.npy', stocks)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:14:21.418961Z",
     "start_time": "2024-03-05T12:14:21.409376Z"
    }
   },
   "id": "bd63e90dbb9c8d10",
   "execution_count": 75
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "np.save(f'{code}/ror.npy', ret)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:14:02.244849Z",
     "start_time": "2024-03-05T12:14:02.233810Z"
    }
   },
   "id": "4051bdb20ae20065",
   "execution_count": 74
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "dict_industry = {\n",
    "            '110000': '农林牧渔',\n",
    "            '220000': '基础化工',\n",
    "            '230000': '钢铁',   \n",
    "            '240000': '有色金属',  \n",
    "            '270000': '电子',  \n",
    "            '280000': '汽车',   \n",
    "            '330000': '家用电器',  \n",
    "            '340000': '食品饮料',   \n",
    "            '350000': '纺织服饰',   \n",
    "            '360000': '轻工制造',   \n",
    "            '370000': '医药生物',   \n",
    "            '410000': '公用事业',   \n",
    "            '420000': '交通运输',   \n",
    "            '430000': '房地产',  \n",
    "            '450000': '商贸零售',  \n",
    "            '460000': '社会服务',   \n",
    "            '480000': '银行',   \n",
    "            '490000': '非银金融',  \n",
    "            '510000': '综合',  \n",
    "            '610000': '建筑材料',  \n",
    "            '620000': '建筑装饰',  \n",
    "            '630000': '电力设备',  \n",
    "            '640000': '机械设备',  \n",
    "            '650000': '国防军工',\n",
    "            '710000': '计算机',  \n",
    "            '720000': '传媒',  \n",
    "            '730000': '通信',  \n",
    "            '740000': '煤炭', \n",
    "            '210000': '煤炭',\n",
    "            '750000': '石油石化',  \n",
    "            '760000': '环保',  \n",
    "            '770000': '美容护理',  \n",
    "        }\n",
    "\n",
    "dict = dict_industry"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:02:13.315238Z",
     "start_time": "2024-03-05T12:02:13.300206Z"
    }
   },
   "id": "d047abcafe0a13c8",
   "execution_count": 43
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "matrix = np.zeros((stock_num, stock_num))\n",
    "for i, stock1 in enumerate(stock_unique):\n",
    "    for j, stock2 in enumerate(stock_unique):\n",
    "        if dict.get(str(unique_mapping.get(stock1))) == dict.get(str(unique_mapping.get(stock2))):\n",
    "            matrix[i][j] = 1\n",
    "np.fill_diagonal(matrix, 1)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:02:13.361Z",
     "start_time": "2024-03-05T12:02:13.317678Z"
    }
   },
   "id": "9c381cfb22e0f04e",
   "execution_count": 44
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "(140, 140)"
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "matrix.shape"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:02:13.376700Z",
     "start_time": "2024-03-05T12:02:13.362Z"
    }
   },
   "id": "9b4dc6edd077fe6f",
   "execution_count": 45
  },
  {
   "cell_type": "code",
   "outputs": [
    {
     "data": {
      "text/plain": "1298"
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.count_nonzero(matrix)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:02:13.391949Z",
     "start_time": "2024-03-05T12:02:13.379945Z"
    }
   },
   "id": "43f144037f6b6c87",
   "execution_count": 46
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "np.save(f'{code}/industry_classification.npy', matrix)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-03-05T12:02:13.407822Z",
     "start_time": "2024-03-05T12:02:13.393949Z"
    }
   },
   "id": "6386e157f30ee615",
   "execution_count": 47
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
